#include "cv.h"

#define COLOR_EQUAL_IMAGE_ELEM( color, im, y, x )						\
	(	CV_IMAGE_ELEM( im, unsigned char, (y), 3*(x)+0 )==color.val[0] &&	\
		CV_IMAGE_ELEM( im, unsigned char, (y), 3*(x)+1 )==color.val[1] &&	\
		CV_IMAGE_ELEM( im, unsigned char, (y), 3*(x)+2 )==color.val[2] )

struct CalcRegionMaskBuffer
{
	CalcRegionMaskBuffer( size_t size ) : maxQueueSize( size )
	{
		queue = new CvPoint[maxQueueSize];
	}
	~CalcRegionMaskBuffer()
	{
		delete [] queue;
	}
	CvPoint* queue;
	const size_t maxQueueSize;
};

void calcRegionMask( IplImage* segIm, CvPoint pt, IplImage* dst, const IplImage* mask, CalcRegionMaskBuffer* buffer )
{
	assert( segIm && dst && mask );
	assert( segIm->nChannels == 3 );
	assert( dst->nChannels == 1 && mask->nChannels == 1 );
	assert( segIm->depth == IPL_DEPTH_8U && dst->depth == IPL_DEPTH_8U && mask->depth == IPL_DEPTH_8U );

	cvZero( dst ); // 0 = unvisited, 254 = in queue, 255 = part of segment
	CvScalar color;
	color.val[0] = CV_IMAGE_ELEM( segIm, unsigned char, pt.y, 3*pt.x+0 );
	color.val[1] = CV_IMAGE_ELEM( segIm, unsigned char, pt.y, 3*pt.x+1 );
	color.val[2] = CV_IMAGE_ELEM( segIm, unsigned char, pt.y, 3*pt.x+2 );

	size_t qsize = 1;
	buffer->queue[0] = pt;

	while( qsize )
	{
		assert( qsize < buffer->maxQueueSize - 10 );
		// remove point from top of the queue
		CvPoint curr = buffer->queue[ --qsize ];
		CV_IMAGE_ELEM(dst, unsigned char, curr.y, curr.x ) = 255; // mark that it has been removed from queue
		
		// check neighbors and add them to queue if same color

		if( curr.y+1 < segIm->height &&  qsize < buffer->maxQueueSize &&
			CV_IMAGE_ELEM(dst, unsigned char, curr.y+1, curr.x )==0 && 
			CV_IMAGE_ELEM(mask, unsigned char, curr.y+1, curr.x )!=0 && 
			COLOR_EQUAL_IMAGE_ELEM( color, segIm, curr.y+1, curr.x ) )
		{
			buffer->queue[ qsize++ ] = cvPoint( curr.x, curr.y+1 );
			CV_IMAGE_ELEM(dst, unsigned char, curr.y+1, curr.x ) = 254; // mark that it is in the queue
		}
		
		if( curr.x > 0 )
		{
			if( curr.y+1 < segIm->height &&  qsize < buffer->maxQueueSize &&
				CV_IMAGE_ELEM(dst, unsigned char, curr.y+1, curr.x-1 )==0 && 
				CV_IMAGE_ELEM(mask, unsigned char, curr.y+1, curr.x-1 )!=0 && 
				COLOR_EQUAL_IMAGE_ELEM( color, segIm, curr.y+1, curr.x-1 ) )
			{
				buffer->queue[ qsize++ ] = cvPoint( curr.x-1, curr.y+1 );
				CV_IMAGE_ELEM(dst, unsigned char, curr.y+1, curr.x-1 ) = 254; // mark that it is in the queue
			}

			if( qsize < buffer->maxQueueSize &&
				CV_IMAGE_ELEM(dst, unsigned char, curr.y, curr.x-1 )==0 && 
				CV_IMAGE_ELEM(mask, unsigned char, curr.y, curr.x-1 )!=0 && 
				COLOR_EQUAL_IMAGE_ELEM( color, segIm, curr.y, curr.x-1 ) )
			{
				buffer->queue[ qsize++ ] = cvPoint( curr.x-1, curr.y );
				CV_IMAGE_ELEM(dst, unsigned char, curr.y, curr.x-1 ) = 254; // mark that it is in the queue
			}
		}

		if( curr.x+1 < segIm->width )
		{

			if( curr.y+1 < segIm->height &&  qsize < buffer->maxQueueSize &&
				CV_IMAGE_ELEM(dst, unsigned char, curr.y+1, curr.x+1 )==0 && 
				CV_IMAGE_ELEM(mask, unsigned char, curr.y+1, curr.x+1 )!=0 && 
				COLOR_EQUAL_IMAGE_ELEM( color, segIm, curr.y+1, curr.x+1 ) )
			{
				buffer->queue[ qsize++ ] = cvPoint( curr.x+1, curr.y+1 );
				CV_IMAGE_ELEM(dst, unsigned char, curr.y+1, curr.x+1 ) = 254; // mark that it is in the queue
			}
			
			if( qsize < buffer->maxQueueSize &&
				CV_IMAGE_ELEM(dst, unsigned char, curr.y, curr.x+1 )==0 && 
				CV_IMAGE_ELEM(mask, unsigned char, curr.y, curr.x+1 )!=0 && 
				COLOR_EQUAL_IMAGE_ELEM( color, segIm, curr.y, curr.x+1 ) )
			{
				buffer->queue[ qsize++ ] = cvPoint( curr.x+1, curr.y );
				CV_IMAGE_ELEM(dst, unsigned char, curr.y, curr.x+1 ) = 254; // mark that it is in the queue
			}
		}
		
		if( curr.y > 0 )
		{
			if( curr.x+1 < segIm->width &&  qsize < buffer->maxQueueSize &&
				CV_IMAGE_ELEM(dst, unsigned char, curr.y-1, curr.x+1 )==0 && 
				CV_IMAGE_ELEM(mask, unsigned char, curr.y-1, curr.x+1 )!=0 && 
				COLOR_EQUAL_IMAGE_ELEM( color, segIm, curr.y-1, curr.x+1 ) )
			{
				buffer->queue[ qsize++ ] = cvPoint( curr.x+1, curr.y-1 );
				CV_IMAGE_ELEM(dst, unsigned char, curr.y-1, curr.x+1 ) = 254; // mark that it is in the queue
			}
			if( curr.x > 0 &&  qsize < buffer->maxQueueSize &&
				CV_IMAGE_ELEM(dst, unsigned char, curr.y-1, curr.x-1 )==0 && 
				CV_IMAGE_ELEM(mask, unsigned char, curr.y-1, curr.x-1 )!=0 && 
				COLOR_EQUAL_IMAGE_ELEM( color, segIm, curr.y-1, curr.x-1 ) )
			{
				buffer->queue[ qsize++ ] = cvPoint( curr.x-1, curr.y-1 );
				CV_IMAGE_ELEM(dst, unsigned char, curr.y-1, curr.x-1 ) = 254; // mark that it is in the queue
			}

			if( qsize < buffer->maxQueueSize &&
				CV_IMAGE_ELEM(dst, unsigned char, curr.y-1, curr.x )==0 && 
				CV_IMAGE_ELEM(mask, unsigned char, curr.y-1, curr.x )!=0 && 
				COLOR_EQUAL_IMAGE_ELEM( color, segIm, curr.y-1, curr.x ) )
			{
				buffer->queue[ qsize++ ] = cvPoint( curr.x, curr.y-1 );
				CV_IMAGE_ELEM(dst, unsigned char, curr.y-1, curr.x ) = 254; // mark that it is in the queue
			}
		}
	}
}